# from tools.screen_recorder import run as run__screen_recorder, ScreenRecorder
from tools.my_reward_analysis import RewardAnalysis
from bdtime import Time
from bdtime import tt
import threading
import os
from tools import my_cv2_utils
import pandas as pd
import numpy as np
from tools.utils import get_public_attrs
from tools import ocr_utils
import re
# from tools.recoder_args import DefaultArgs
from tqdm import tqdm
import cv2
import json
from bdtime import show_ls


def is_fraction(s):
    pattern = r'^\d+/\d+$'
    match = re.match(pattern, s)
    return bool(match)


class Coordinates:
    is_complete_coordinates = [(210, 335), (260, 350)]
    extinct_enemy_coordinates = [(65, 104), (92, 115)]


# class RewardAnalysis:
#     def __init__(self, operation_name, img_dir_name,
#                  record_dir_path=DefaultArgs.record_path.recorded_state_dir_path, i_range: tuple = None):
#         self.record_dir_path = record_dir_path
#         self.screen_dir_path = os.path.join(record_dir_path, 'screen_image')
#
#         img_dir_path = os.path.join(self.screen_dir_path, img_dir_name)
#         assert os.path.exists(img_dir_path)
#         self.img_dir_path = img_dir_path
#
#         key_file_path = os.path.join(record_dir_path, 'key', operation_name)
#         assert os.path.exists(key_file_path)
#         self.key_file_path = key_file_path
#
#         mouse_file_path = os.path.join(record_dir_path, 'rel_mouse', operation_name)
#         assert os.path.exists(mouse_file_path)
#         self.mouse_file_path = mouse_file_path
#
#         with open(key_file_path, 'r') as f:
#             key_ls = json.load(f)
#         key_ls = [i for i in key_ls if 'i' in i]
#
#         with open(mouse_file_path, 'r') as f:
#             mouse_ls = json.load(f)
#         mouse_ls = [i for i in mouse_ls if 'i' in i]
#
#         self.analysis_result_save_dir = f'../tempdir/analysis_result/{img_dir_name}'
#         os.makedirs(self.analysis_result_save_dir, exist_ok=True)
#         self.output_path__reward_state_info_df = os.path.join(self.analysis_result_save_dir, 'reward_state_info_df.csv')
#
#         self._reward_state_info = []
#         self.reward_state_info_df = pd.DataFrame(self._reward_state_info, columns=['game_state', 'extinct_enemy', 'remain_enemy'])
#
#         _img_ls = os.listdir(img_dir_path)
#         _img_ls = [i for i in _img_ls if i.endswith('jpg') or i.endswith('png')]
#         get_create_time = lambda filename: os.path.getctime(os.path.join(img_dir_path, filename))
#         ordered_img_ls = sorted(_img_ls, key=get_create_time)
#
#         if isinstance(i_range, tuple):
#             assert len(i_range) == 2, 'i_range的长度必须为2!'
#             ordered_img_ls = ordered_img_ls[i_range[0]: i_range[1]]
#             key_ls = key_ls[i_range[0]: i_range[1]]
#             mouse_ls = mouse_ls[i_range[0]: i_range[1]]
#
#             self.output_path__reward_state_info_df = os.path.join(self.analysis_result_save_dir, f'reward_state_info_df__{i_range[0]}_to_{i_range[1]}.csv')
#
#         img_path_ls = [os.path.join(img_dir_path, i) for i in ordered_img_ls]
#         info_ls = [os.path.splitext(i)[0].split('__')[1:] for i in ordered_img_ls]
#
#         dtype = {'i': np.int32, 'time': np.float32}
#         df = pd.DataFrame(info_ls, columns=['i', 'time']).astype(dtype)
#         col = pd.Series(ordered_img_ls)
#         df['file_name'] = col
#         col = pd.Series(img_path_ls)
#         df['file_path'] = col
#
#         col = pd.Series(mouse_ls)
#         df['mouse'] = col
#         col = pd.Series(key_ls)
#         df['key'] = col
#
#         # show_ls(key_ls[:5])
#         # show_ls(mouse_ls[:5])
#
#         # len(key_ls), len(mouse_ls), len(img_path_ls)
#
#         self.key_ls = key_ls
#         self.mouse_ls = mouse_ls
#
#         self.df = df
#         self.img_path_ls = img_path_ls
#         self.ordered_img_ls = ordered_img_ls
#         self.info_ls = info_ls
#
#         self.total_frames = len(img_path_ls)
#         self.key_total_frames = len(key_ls)
#
#         self.coordinates = Coordinates()
#
#         self.tt = Time()
#
#         self.cache_info_dc = {}
#
#         self.total_enemy = -1
#         self.get_total_enemy()
#
#     def get_mid_val(self, game_state, extinct_enemy):
#         if game_state == 0:
#             mid_val = extinct_enemy
#         elif game_state == -1:
#             mid_val = -1
#         elif game_state == 1:
#             mid_val = self.total_enemy
#         else:
#             raise ValueError('game_state取值错误! 必须在[-1, 0, 1]之中!')
#         return mid_val
#
#     def get_img_i(self, i):
#         img_i = my_cv2_utils.read_image(self.img_path_ls[i])
#         return img_i
#
#     def show_img(self, img_i):
#         if isinstance(img_i, int):
#             img_i = self.get_img_i(img_i)
#         my_cv2_utils.show_img(img_i)
#
#     def get_info(self, img_i, show=False):
#         _img_i = None
#         if isinstance(img_i, int):
#             _img_i = img_i
#             if img_i in self.cache_info_dc:
#                 game_state, extinct_enemy, total_enemy = self.cache_info_dc.get(img_i)
#                 return game_state, extinct_enemy, total_enemy
#
#             img_i = self.get_img_i(img_i)
#
#         extinct_enemy, total_enemy = self._get_extinct_enemy_and_total_enemy(img_i)
#
#         is_complete = self._get_is_complete(img_i)
#         is_start = not is_complete and extinct_enemy is None and total_enemy is None
#
#         game_state = 0
#         if is_start:
#             game_state = -1
#         elif is_complete:
#             game_state = 1
#
#         if show:
#             self.show_img(img_i)
#
#         if _img_i is not None:
#             self.cache_info_dc.update({_img_i: (game_state, extinct_enemy, total_enemy)})
#
#         return game_state, extinct_enemy, total_enemy
#
#     def _get_is_complete(self, img_i, show=False):
#         if isinstance(img_i, int):
#             img_i = self.get_img_i(img_i)
#         text, score = ocr_utils.get_orc_result(img_i, coordinates=self.coordinates.is_complete_coordinates, show=show)
#         res = '已' in text
#         return res
#
#     def _get_extinct_enemy_and_total_enemy(self, img_i, show=False):
#         if isinstance(img_i, int):
#             img_i = self.get_img_i(img_i)
#         text, score = ocr_utils.get_orc_result(img_i, coordinates=self.coordinates.extinct_enemy_coordinates, show=show)
#
#         if is_fraction(text):
#             extinct_enemy, total_enemy = [int(i) for i in text.split('/')]
#         else:
#             extinct_enemy, total_enemy = None, None
#         return extinct_enemy, total_enemy
#
#     def show_base_df(self, show_path=False):
#         columns = self.df.columns if show_path else ['i', 'time', 'file_name']
#         print(self.df[columns].head())
#
#     def get_total_enemy(self, debug=False):
#         self.total_enemy = self._get_total_enemy_by_binary_search(debug=debug)
#         assert self.total_enemy != -1
#
#     def _get_total_enemy_by_binary_search(self, debug=False):
#
#         """
#         这里要找到[game_state, extinct_enemy, total_enemy]中的`total_enemy`, 即一个game_state为0的值
#         """
#         # target = 1
#         # arr = self.img_path_ls
#
#         length = self.total_frames
#
#         low, high = 0, length - 1
#
#         target = 0
#         while low <= high:
#             mid = (low + high) // 2
#             game_state, extinct_enemy, total_enemy = self.get_info(mid)
#             mid_val = game_state
#             if debug:
#                 print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}')
#             if mid_val == target:
#                 if debug:
#                     print(f'~~~~~~ mid_val == target --- mid: {mid}, mid_val: {mid_val}')
#                 return total_enemy  # 找到目标，返回索引
#             elif mid_val < target:
#                 low = mid + 1  # 目标在右半部分
#             else:
#                 high = mid - 1  # 目标在左半部分
#
#         return -1
#
#     def binary_search(self, target_v=0, change_dc=None, range_dc=None):
#         """
#         这里要找到[game_state, extinct_enemy, total_enemy]中的`total_enemy`, 即一个game_state为0的值
#         """
#         # target = 1
#         # arr = self.img_path_ls
#
#         length = self.total_frames
#
#         self.total_enemy
#
#         if change_dc is None:
#             # change_ls = list(range(-1, self.total_enemy + 1))  # start: -1 -> 0, end: 8 -> 9
#             change_ls = list(range(self.total_enemy + 1))  # start: -1 -> 0, end: 8 -> 9
#             change_dc = dict(zip(change_ls, [-1] * len(change_ls)))
#             range_dc = change_dc.copy()
#
#         # range_dc[0] = (0, 100)
#
#         sum(change_dc.values()) == -len(change_ls)  # 全为-1
#         list(change_dc.values()).count(-1) == 0  # 不存在-1
#
#         min_i, max_i = 1, length - 1
#
#         change_dc = dict(zip(change_ls, [-1] * len(change_ls)))
#         range_dc = change_dc.copy()
#         target_dc = change_dc.copy()
#
#         for target_v in change_ls:
#         # for target_v in [0, 1, 2]:
#         # for target_v in [7]:
#             low, high = 1, length - 1
#             # target_v = 3
#             target_i = -1
#             # mid = -1
#             # mid_val = -1
#             tt.__init__()
#             # while low <= high and tt.during(10):
#             while low <= high:
#                 mid = (low + high) // 2
#
#                 game_state, extinct_enemy, _ = self.get_info(mid)
#                 mid_val = self.get_mid_val(game_state, extinct_enemy)
#
#                 _game_state, _extinct_enemy, _ = self.get_info(mid - 1)
#                 _mid_val = self.get_mid_val(_game_state, _extinct_enemy)
#
#                 print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}')
#
#                 if mid_val == target_v:
#                     print(f'~~~~~~ mid: {mid}, mid_val: {mid_val}')
#                     target_i = mid
#
#                     # high = mid - 1  # 目标在左半部分
#
#                     i = mid_val
#                     v = mid
#                     if range_dc[i] == -1:
#                         range_dc[i] = [min_i, v]
#                     else:
#                         range_dc[i][1] = min(v, range_dc[i][1])
#
#                     # i = mid
#                     # range_dc[i] =
#                     # if range_dc[i] == -1:
#                     #     range_dc[i] = [low, max_i]
#                     # else:
#                     #     if low > range_dc[i][0]:
#                     #         range_dc[i][0] = low
#                     break
#                     # return mid  # 找到目标，返回索引
#                 elif mid_val < target_v:
#                     low = mid + 1  # 目标在右半部分
#                     # change_dc
#                     # range_dc
#                     # range_dc = change_dc.copy()
#
#                     for i in range(mid_val + 1, self.total_enemy + 1):
#                         if range_dc[i] == -1:
#                             range_dc[i] = [low, max_i]
#                         else:
#                             range_dc[i][0] = max(low, range_dc[i][0])
#                             # if low > range_dc[i][0]:
#                             #     range_dc[i][0] = low
#                 else:
#                     high = mid - 1  # 目标在左半部分
#
#                     for i in range(0, mid_val):
#                         if range_dc[i] == -1:
#                             range_dc[i] = [min_i, high]
#                         else:
#                             range_dc[i][1] = min(high, range_dc[i][1])
#                             # if high < range_dc[i][1]:
#                             #     range_dc[i][1] = high
#                 print(f'--- target_v: {target_v} ~~~~~~ range_dc: {range_dc} ====== now: {tt.now()}')
#
#                 # print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}, target_v: {target_v}')
#
#                 # mid_val_0, _, _ = self.get_info(mid)
#                 # mid_val_1, _, _ = self.get_info(mid + 1)
#                 # if mid_val_0 == mid_val_1:
#                 #     if mid_val_0 == 0:
#
#             # assert mid_val == target
#             # assert target_i != -1
#
#             if target_i != -1:
#                 change_dc[target_v] = target_i
#                 target_dc[target_v] = target_i
#
#             print(f'====== target_v: {target_v}--- range_dc: {range_dc} ====== now: {tt.now()}\n')
#
#         sum([v[1] - v[0] for k, v in range_dc.items()])
#         return target_i  # -1则目标不在列表中
#
#         change_dc
#         range_dc
#         self.get_info(target_i)
#
#         self.cache_info_dc
#
#         index = 659
#         self.get_info(index)
#         self.get_mid_val(*self.get_info(index)[:2])
#
#         self.get_info(index, show=True)
#         self._get_is_complete(index, show=True)
#
#         return target_i  # -1则目标不在列表中
#
#     def get_reward_per_frame(self, batch_size=1, get_rtg=True):
#         _reward_state_info = []
#         tq_i = tqdm(total=self.total_frames)
#
#         for i in range(0, self.total_frames, batch_size):
#             # print(f'i --- {i}')
#             game_state, extinct_enemy, _ = self.get_info(i)
#             extinct_enemy = extinct_enemy if extinct_enemy else self.total_enemy if game_state == 1 else 0
#             _reward_state_info.append([game_state, extinct_enemy, self.total_enemy - extinct_enemy])
#             tq_i.update(batch_size)
#
#         self._reward_state_info = _reward_state_info
#         df = pd.DataFrame(self._reward_state_info, columns=['game_state', 'extinct_enemy', 'remain_enemy'])
#         if get_rtg:
#             df = self.get_rtg(df)
#
#         merged_df = pd.concat([self.df, df], axis=1)
#         self.reward_state_info_df = merged_df
#         self.reward_state_info_df.to_csv(self.output_path__reward_state_info_df, index=False)
#
#     def get_rtg(self, df, delta__reward_time=None, delta__reward_extinct_enemy=None, delta__reward_change_game_state=None):
#         if delta__reward_time is None:
#             # delta__reward_time = -1  # 每一帧扣1分
#             delta__reward_time = -1 / DefaultArgs.frames  # 每一秒扣1分
#             # delta__reward_time = -1 / DefaultArgs.frames / 60  # 每一分钟扣1分
#
#         if delta__reward_extinct_enemy is None:
#             delta__reward_extinct_enemy = 1
#             # delta__reward_extinct_enemy = 100 / self.total_enemy
#
#         if delta__reward_change_game_state is None:
#             delta__reward_change_game_state = 10
#             # delta__reward_change_game_state = 100
#
#         df['reward'] = delta__reward_time
#
#         last_extinct_enemy = 0
#         last_game_state = -1
#         for i, row in df.iterrows():
#             # break
#             extinct_enemy = row['extinct_enemy']
#             game_state = row['game_state']
#             reward = row['reward']
#             # extinct_enemy = df.loc[i, 'extinct_enemy']
#             # game_state = df.loc[i, 'game_state']
#             # reward = 0
#             if extinct_enemy > last_extinct_enemy:
#                 reward += delta__reward_extinct_enemy
#                 print(f'--- extinct_enemy i: {i}')
#             if game_state > last_game_state:
#                 reward += delta__reward_change_game_state
#                 print(f'--- game_state i: {i}')
#             df.loc[i, 'reward'] = reward
#
#             if i < df.shape[0] - 1:
#                 last_extinct_enemy = extinct_enemy
#                 last_game_state = game_state
#
#         reward_ls = df['reward'].tolist()
#
#         rtg_ls = []
#         rtg_accumulate = 0
#         for i in reversed(range(len(reward_ls))):
#             # print(i)
#             reward_i = reward_ls[i]
#
#             rtg_accumulate += reward_i
#             rtg_ls.append(rtg_accumulate)
#         rtg_ls.reverse()
#
#         df['rtg'] = rtg_ls
#         return df
#
#     def show_video(self, start_i, end_i):
#         image_files = self.img_path_ls[start_i: end_i]
#         total = end_i - start_i
#         tq_i = tqdm(total=total)
#
#         record_time_schedule = DefaultArgs.record_time_schedule[start_i: end_i]
#
#         tt = Time()
#         for i in range(len(image_files)):
#             image_file = image_files[i]
#
#             next_record_time = record_time_schedule[i] - record_time_schedule[0]
#
#             now = tt.now()
#             while now < next_record_time:
#                 if DefaultArgs.use_sleep_wait:
#                     tt.sleep(0.001)
#                 now = tt.now()
#
#             # msg = f'now: {now}, next_record_time: {next_record_time: .3f}'
#             msg = f'delta_t: {now - next_record_time: .3f}'
#             print(msg)
#
#             # 读取图片
#             # frame = cv2.imread(image_file)
#             frame = my_cv2_utils.read_image(image_file)
#             # my_cv2_utils.show_img(frame)
#
#             # 显示图片
#             cv2.imshow('Frame', frame)
#
#             # 检测键盘按键，如果是 'q' 则退出循环
#             if cv2.waitKey(25) & 0xFF == ord('q'):
#                 break
#
#             tq_i.update(1)
#             # tt.sleep(interval)
#
#         # 关闭窗口
#         cv2.destroyAllWindows()

# from tools.my_reward_analysis import RewardAnalysis
# from tools import RewardAnalysis


def main():
    # img_dir_path = '../tempdir/recorded_states/screen_image/test_7'
    # img_dir_name = 'test_7'
    # operation_name = 'test.json'

    img_dir_name = 'cao_shen.json__2024-02-02 08-16-52'
    operation_name = 'cao_shen.json'
    record_analysis = RewardAnalysis(operation_name=operation_name, img_dir_name=img_dir_name)

    # record_analysis.df.columns
    # df_2 = pd.read_csv(record_analysis.output_path__reward_state_info_df)
    # df_2.head()
    # record_analysis.df.head()
    # record_analysis = RewardAnalysis(operation_name=operation_name, img_dir_name=img_dir_name, i_range=(240, 260))
    # record_analysis.total_frames

    # print(record_analysis.img_dir_path)

    # record_analysis.get_total_enemy(debug=True)
    # record_analysis.total_enemy

    index = record_analysis.key_total_frames

    # duration = 10
    # index = record_analysis.key_total_frames - duration
    # # record_analysis.show_video(index, index + duration)
    # # Coordinates.extinct_enemy_coordinates
    # # record_analysis._get_extinct_enemy_and_total_enemy(260, show=True)
    #
    # game_state, extinct_enemy, total_enemy = record_analysis.get_info(index)
    # print(f"--- index: {index} --- game_state: {game_state} | extinct_enemy, total_enemy: {extinct_enemy, total_enemy}")

    # is_complete = record_analysis._get_is_complete(index)
    # ocr_utils.get_orc_result(record_analysis.get_img_i(index), coordinates=record_analysis.coordinates.is_complete_coordinates)
    # ocr_utils.get_orc_result(record_analysis.get_img_i(index), coordinates=record_analysis.coordinates.extinct_enemy_coordinates)

    # record_analysis.show_img(index)

    record_analysis.get_reward_per_frame()

    pass


if __name__ == '__main__':
    main()






